import numpy as np
import matplotlib.pyplot as plt
import time
from qp_taichi import QPSolver

def mpc_control_problem(A, B, Q, R, N, x0, xref, umin, umax, xmin, xmax):
    """
    构建MPC控制问题的QP表示
    
    :param A: 系统动态矩阵
    :param B: 输入矩阵
    :param Q: 状态成本矩阵
    :param R: 控制成本矩阵
    :param N: 预测步长
    :param x0: 初始状态
    :param xref: 参考状态轨迹
    :param umin: 控制输入下限
    :param umax: 控制输入上限
    :param xmin: 状态下限
    :param xmax: 状态上限
    :return: P, q, G, h, A_eq, b_eq - QP问题参数
    """
    # 问题维度
    nx = A.shape[0]  # 状态维度
    nu = B.shape[1]  # 控制输入维度
    
    # 决策变量维度 (状态和控制输入)
    nz = (N + 1) * nx + N * nu
    
    # 构建二次目标函数 1/2 z^T P z + q^T z
    P = np.zeros((nz, nz))
    q = np.zeros(nz)
    
    # 填充成本矩阵 (状态成本)
    for i in range(N + 1):
        idx = i * nx
        P[idx:idx+nx, idx:idx+nx] = Q
        q[idx:idx+nx] = -Q @ xref[i]
    
    # 最后一个状态的额外终端成本
    P[N*nx:(N+1)*nx, N*nx:(N+1)*nx] = Q * 10
    
    # 填充成本矩阵 (控制输入成本)
    for i in range(N):
        idx = (N + 1) * nx + i * nu
        P[idx:idx+nu, idx:idx+nu] = R
    
    # 系统动态约束 x_{k+1} = A x_k + B u_k
    A_dyn = np.zeros((N * nx, nz))
    b_dyn = np.zeros(N * nx)
    
    for i in range(N):
        # x_k
        idx_x = i * nx
        A_dyn[i*nx:(i+1)*nx, idx_x:idx_x+nx] = A
        
        # x_{k+1}
        idx_x_next = (i+1) * nx
        A_dyn[i*nx:(i+1)*nx, idx_x_next:idx_x_next+nx] = -np.eye(nx)
        
        # u_k
        idx_u = (N + 1) * nx + i * nu
        A_dyn[i*nx:(i+1)*nx, idx_u:idx_u+nu] = B
    
    # 初始状态约束 x_0 = x0
    A_init = np.zeros((nx, nz))
    A_init[:, :nx] = np.eye(nx)
    b_init = x0
    
    # 合并等式约束
    A_eq = np.vstack([A_dyn, A_init])
    b_eq = np.concatenate([b_dyn, b_init])
    
    # 不等式约束
    # 控制输入约束 umin <= u <= umax
    G_u_lower = np.zeros((N * nu, nz))
    G_u_upper = np.zeros((N * nu, nz))
    h_u_lower = np.zeros(N * nu)
    h_u_upper = np.zeros(N * nu)
    
    for i in range(N):
        idx_u = (N + 1) * nx + i * nu
        G_u_lower[i*nu:(i+1)*nu, idx_u:idx_u+nu] = -np.eye(nu)
        G_u_upper[i*nu:(i+1)*nu, idx_u:idx_u+nu] = np.eye(nu)
        h_u_lower[i*nu:(i+1)*nu] = -umin
        h_u_upper[i*nu:(i+1)*nu] = umax
    
    # 状态约束 xmin <= x <= xmax
    G_x_lower = np.zeros(((N+1) * nx, nz))
    G_x_upper = np.zeros(((N+1) * nx, nz))
    h_x_lower = np.zeros((N+1) * nx)
    h_x_upper = np.zeros((N+1) * nx)
    
    for i in range(N+1):
        idx_x = i * nx
        G_x_lower[i*nx:(i+1)*nx, idx_x:idx_x+nx] = -np.eye(nx)
        G_x_upper[i*nx:(i+1)*nx, idx_x:idx_x+nx] = np.eye(nx)
        h_x_lower[i*nx:(i+1)*nx] = -xmin
        h_x_upper[i*nx:(i+1)*nx] = xmax
    
    # 合并不等式约束
    G = np.vstack([G_u_lower, G_u_upper, G_x_lower, G_x_upper])
    h = np.concatenate([h_u_lower, h_u_upper, h_x_lower, h_x_upper])
    
    return P, q, G, h, A_eq, b_eq

def double_integrator_system():
    """
    创建一个双积分器系统
    表示一个简单的二阶系统，如位置-速度控制
    
    :return: A, B - 系统动态和输入矩阵
    """
    # 离散化时间步长
    dt = 0.1
    
    # 连续时间系统矩阵
    Ac = np.array([
        [0, 1],
        [0, 0]
    ])
    
    Bc = np.array([
        [0],
        [1]
    ])
    
    # 离散化
    A = np.eye(2) + dt * Ac
    B = dt * Bc
    
    return A, B

def main():
    """
    使用QP-Taichi解决MPC控制问题的示例
    """
    print("QP-Taichi MPC控制示例")
    
    # 系统参数 (双积分器)
    A, B = double_integrator_system()
    nx = A.shape[0]  # 状态维度
    nu = B.shape[1]  # 控制输入维度
    
    # MPC参数
    N = 20  # 预测步长
    
    # 成本矩阵
    Q = np.diag([10.0, 1.0])  # 状态成本 (位置权重大，速度权重小)
    R = np.diag([0.1])        # 控制成本
    
    # 约束
    umin = np.array([-2.0])  # 控制输入下限
    umax = np.array([2.0])   # 控制输入上限
    xmin = np.array([-10.0, -5.0])  # 状态下限
    xmax = np.array([10.0, 5.0])    # 状态上限
    
    # 初始状态 [位置, 速度]
    x0 = np.array([5.0, 0.0])
    
    # 参考轨迹 (目标是原点)
    xref = [np.zeros(nx) for _ in range(N+1)]
    
    # 创建QP求解器
    try:
        solver = QPSolver(arch="gpu")
    except Exception as e:
        print(f"GPU初始化失败: {e}")
        print("使用CPU代替...")
        solver = QPSolver(arch="cpu")
    
    # 仿真参数
    sim_time = 10.0  # 仿真时间
    dt = 0.1        # 采样时间
    n_steps = int(sim_time / dt)
    
    # 仿真变量
    states = np.zeros((n_steps+1, nx))
    states[0] = x0
    controls = np.zeros((n_steps, nu))
    solve_times = np.zeros(n_steps)
    
    print("开始MPC仿真...")
    for step in range(n_steps):
        # 当前状态
        x_cur = states[step]
        
        # 构建MPC问题
        t_start = time.time()
        P, q, G, h, A_eq, b_eq = mpc_control_problem(
            A, B, Q, R, N, x_cur, xref, umin, umax, xmin, xmax
        )
        
        # 求解QP问题
        result = solver.solve(P, q, G, h, A_eq, b_eq,
                             abs_tol=1e-3, rel_tol=1e-2,
                             max_iter=1000, verbose=False)
        t_end = time.time()
        solve_times[step] = t_end - t_start
        
        # 提取控制输入 (仅使用第一个控制输入)
        z_opt = result.x
        u = z_opt[(N+1)*nx:(N+1)*nx+nu]
        controls[step] = u
        
        # 系统状态更新
        states[step+1] = A @ states[step] + B @ u
        
        # 打印进度
        if step % 10 == 0:
            print(f"步骤 {step}/{n_steps}, 状态: {states[step+1]}, " 
                  f"控制: {u}, 求解时间: {solve_times[step]:.3f}秒")
    
    # 绘制结果
    plt.figure(figsize=(15, 10))
    
    # 位置变化
    plt.subplot(4, 1, 1)
    plt.plot(np.arange(0, sim_time+dt, dt), states[:, 0])
    plt.axhline(y=xmin[0], color='r', linestyle='--', alpha=0.5)
    plt.axhline(y=xmax[0], color='r', linestyle='--', alpha=0.5)
    plt.grid(True)
    plt.ylabel('位置')
    plt.title('MPC控制 - 双积分器系统')
    
    # 速度变化
    plt.subplot(4, 1, 2)
    plt.plot(np.arange(0, sim_time+dt, dt)[:-1], states[:-1, 1])
    plt.axhline(y=xmin[1], color='r', linestyle='--', alpha=0.5)
    plt.axhline(y=xmax[1], color='r', linestyle='--', alpha=0.5)
    plt.grid(True)
    plt.ylabel('速度')
    
    # 控制输入
    plt.subplot(4, 1, 3)
    plt.step(np.arange(0, sim_time, dt), controls)
    plt.axhline(y=umin[0], color='r', linestyle='--', alpha=0.5)
    plt.axhline(y=umax[0], color='r', linestyle='--', alpha=0.5)
    plt.grid(True)
    plt.ylabel('控制输入')
    
    # 求解时间
    plt.subplot(4, 1, 4)
    plt.plot(np.arange(0, sim_time, dt), solve_times * 1000)  # 转换为毫秒
    plt.grid(True)
    plt.xlabel('时间 (秒)')
    plt.ylabel('求解时间 (毫秒)')
    
    plt.tight_layout()
    plt.savefig('mpc_results.png')
    
    print("\nMPC仿真完成")
    print(f"平均求解时间: {np.mean(solve_times) * 1000:.2f} 毫秒")
    print(f"最终状态: {states[-1]}")
    print("结果图已保存到 mpc_results.png")

if __name__ == "__main__":
    main() 